import pymysql
import datetime
from werkzeug.security import generate_password_hash

def get_db_connection():
    """创建数据库连接"""
    connection = pymysql.connect(
        host='localhost',
        port=3306,
        user='root',
        password='123456',
        database='plate'
    )
    return connection

def init_db():
    """初始化数据库：创建表、清空数据、添加测试数据"""
    conn = get_db_connection()
    cursor = conn.cursor()
    
    print("开始初始化数据库...")
    
    # 创建表（如果不存在）
    print("创建表结构...")
    
    # 用户表
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS users (
        id INT AUTO_INCREMENT PRIMARY KEY,
        username VARCHAR(100) NOT NULL,
        password VARCHAR(255) NOT NULL,
        phone VARCHAR(20) NOT NULL UNIQUE,
        role VARCHAR(20) NOT NULL,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    ''')
    
    # 车辆表
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS vehicles (
        id INT AUTO_INCREMENT PRIMARY KEY,
        plate_number VARCHAR(20) NOT NULL UNIQUE,
        owner_id INT NOT NULL,
        vehicle_type VARCHAR(50),
        color VARCHAR(30),
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        FOREIGN KEY (owner_id) REFERENCES users(id)
    )
    ''')
    
    # 停车场表
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS parking_lots (
        id INT AUTO_INCREMENT PRIMARY KEY,
        name VARCHAR(100) NOT NULL,
        address VARCHAR(255),
        total_spaces INT NOT NULL,
        hourly_rate DECIMAL(10,2) NOT NULL,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    ''')
    
    # 停车记录表
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS parking_records (
        id INT AUTO_INCREMENT PRIMARY KEY,
        vehicle_id INT NOT NULL,
        plate_number VARCHAR(20) NOT NULL,
        parking_lot_id INT NOT NULL,
        entry_time DATETIME NOT NULL,
        exit_time DATETIME,
        fee DECIMAL(10,2),
        payment_status VARCHAR(20) DEFAULT 'unpaid',
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        FOREIGN KEY (vehicle_id) REFERENCES vehicles(id),
        FOREIGN KEY (parking_lot_id) REFERENCES parking_lots(id)
    )
    ''')
    
    # 停车位预约表
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS parking_reservations (
        id INT AUTO_INCREMENT PRIMARY KEY,
        user_id INT NOT NULL,
        vehicle_id INT NOT NULL,
        plate_number VARCHAR(20) NOT NULL,
        parking_lot_id INT NOT NULL,
        reservation_time DATETIME NOT NULL,
        expiration_time DATETIME NOT NULL,
        status VARCHAR(20) DEFAULT 'active',
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        FOREIGN KEY (user_id) REFERENCES users(id),
        FOREIGN KEY (vehicle_id) REFERENCES vehicles(id),
        FOREIGN KEY (parking_lot_id) REFERENCES parking_lots(id)
    )
    ''')
    
    # 清空所有表数据
    print("清空表数据...")
    # 按照外键约束顺序删除
    cursor.execute("SET FOREIGN_KEY_CHECKS = 0")
    cursor.execute("TRUNCATE TABLE parking_reservations")
    cursor.execute("TRUNCATE TABLE parking_records")
    cursor.execute("TRUNCATE TABLE vehicles")
    cursor.execute("TRUNCATE TABLE parking_lots")
    cursor.execute("TRUNCATE TABLE users")
    cursor.execute("SET FOREIGN_KEY_CHECKS = 1")
    
    # 添加测试数据
    print("添加测试数据...")
    
    # 插入停车场数据
    print("1. 添加停车场...")
    cursor.execute('''
    INSERT INTO parking_lots (name, address, total_spaces, hourly_rate) 
    VALUES 
    ('A区停车场', '市中心广场东侧', 10, 10.00),
    ('B区停车场', '市中心广场西侧', 10, 8.00)
    ''')
    
    # 插入管理员账号
    print("2. 添加管理员账号...")
    admin_password = generate_password_hash('admin123')
    cursor.execute('''
    INSERT INTO users (username, password, phone, role) 
    VALUES ('管理员', %s, '13800000000', 'admin')
    ''', (admin_password,))
    
    # 添加测试用户数据
    print("3. 添加测试用户和车辆...")
    test_password = generate_password_hash('123456')
    cursor.execute('''
    INSERT INTO users (username, password, phone, role) 
    VALUES ('测试用户', %s, '13900000001', 'user')
    ''', (test_password,))
    
    # 获取测试用户ID
    cursor.execute("SELECT id FROM users WHERE username='测试用户'")
    test_user_id = cursor.fetchone()[0]
    
    # 为测试用户添加车辆
    cursor.execute('''
    INSERT INTO vehicles (plate_number, owner_id, vehicle_type, color)
    VALUES ('京A12345', %s, '小型轿车', '黑色')
    ''', (test_user_id,))
    
    # 获取测试用户车辆ID
    cursor.execute("SELECT id FROM vehicles WHERE plate_number='京A12345'")
    test_vehicle_id = cursor.fetchone()[0]
    
    # 添加10个模拟用户及其车辆
    print("4. 添加模拟用户和车辆...")
    for i in range(1, 11):
        # 创建用户
        sim_password = generate_password_hash('123456')
        phone = f'138{i:08d}'
        cursor.execute('''
        INSERT INTO users (username, password, phone, role) 
        VALUES (%s, %s, %s, 'user')
        ''', (f'模拟用户{i}', sim_password, phone))
        
        # 获取用户ID
        cursor.execute(f"SELECT id FROM users WHERE phone='{phone}'")
        sim_user_id = cursor.fetchone()[0]
        
        # 添加车辆
        plate = f'京B{i:05d}'
        cursor.execute('''
        INSERT INTO vehicles (plate_number, owner_id, vehicle_type, color)
        VALUES (%s, %s, '小型轿车', '白色')
        ''', (plate, sim_user_id))
        
        # 获取车辆ID
        cursor.execute(f"SELECT id FROM vehicles WHERE plate_number='{plate}'")
        sim_vehicle_id = cursor.fetchone()[0]
        
        # 为A区停车场添加10个占用记录（满车）
        if i <= 10:
            entry_time = datetime.datetime.now() - datetime.timedelta(hours=i)
            cursor.execute('''
            INSERT INTO parking_records (vehicle_id, plate_number, parking_lot_id, entry_time)
            VALUES (%s, %s, 1, %s)
            ''', (sim_vehicle_id, plate, entry_time))
        
        # 为B区停车场添加5个占用记录（半满）
        if i <= 5:
            entry_time = datetime.datetime.now() - datetime.timedelta(hours=i)
            cursor.execute('''
            INSERT INTO parking_records (vehicle_id, plate_number, parking_lot_id, entry_time)
            VALUES (%s, %s, 2, %s)
            ''', (sim_vehicle_id, plate, entry_time))
    
    conn.commit()
    cursor.close()
    conn.close()
    
    print("数据库初始化完成!")

if __name__ == "__main__":
    """直接运行此脚本时执行初始化"""
    init_db() 